import os
import json
import numpy as np
import joblib
import jieba
import pickle
from sklearn.feature_extraction.text import TfidfVectorizer

# 定义分词函数，用于TF-IDF向量化器
def jieba_tokenizer(text):
    """
    使用jieba分词器对文本进行分词
    
    Args:
        text: 输入文本
        
    Returns:
        分词结果列表
    """
    return list(jieba.cut(text, cut_all=False))

class TechStackPredictor:
    def __init__(self, model_dir="../../../model/tech_stack_model", threshold=0.3):
        """
        技术栈预测器，用于预测面试题所涉及的技术栈
        
        Args:
            model_dir: 模型目录，默认为项目根目录下的model/tech_stack_model
            threshold: 分类阈值，默认0.3，用于确定是否将某个技术栈包含在结果中
        """
        self.model_dir = model_dir
        self.threshold = threshold
        
        # 加载技术栈列表
        tech_stacks_path = os.path.join(model_dir, "tech_stacks.json")
        if os.path.exists(tech_stacks_path):
            with open(tech_stacks_path, "r", encoding="utf-8") as f:
                self.tech_stacks = json.load(f)
        else:
            raise FileNotFoundError(f"技术栈列表文件不存在: {tech_stacks_path}")
        
        # 加载向量化器
        vectorizer_path = os.path.join(model_dir, "vectorizer.pkl")
        if os.path.exists(vectorizer_path):
            self.vectorizer = joblib.load(vectorizer_path)
        else:
            raise FileNotFoundError(f"向量化器文件不存在: {vectorizer_path}")
        
        # 加载模型
        model_path = os.path.join(model_dir, "tech_stack_model.pkl")
        if os.path.exists(model_path):
            self.model = joblib.load(model_path)
        else:
            raise FileNotFoundError(f"模型文件不存在: {model_path}")
        
        print(f"技术栈预测器已初始化")
    
    def predict(self, question, return_scores=False):
        """
        预测问题的技术栈
        
        Args:
            question: 问题文本
            return_scores: 是否返回分数，默认为False
        
        Returns:
            如果return_scores为False，返回预测的技术栈列表
            如果return_scores为True，返回预测的技术栈列表和每个技术栈的置信度分数
        """
        # 特征提取：将问题文本转换为TF-IDF特征
        X = self.vectorizer.transform([question])
        
        # 预测：获取每个标签为正类的概率
        y_pred_proba = self.model.predict_proba(X)
        scores = np.array([p[0, 1] for p in y_pred_proba])
        
        # 根据阈值筛选出预测的技术栈
        pred_indices = np.where(scores > self.threshold)[0]
        predicted_techs = [self.tech_stacks[i] for i in pred_indices]
        
        # 如果没有预测结果，返回分数最高的技术栈
        if not predicted_techs:
            top_idx = np.argmax(scores)
            predicted_techs = [self.tech_stacks[top_idx]]
        
        if return_scores:
            # 返回每个技术栈的置信度分数
            tech_scores = {self.tech_stacks[i]: float(scores[i]) for i in range(len(self.tech_stacks))}
            return predicted_techs, tech_scores
        else:
            return predicted_techs
    
    def batch_predict(self, questions, return_scores=False):
        """
        批量预测问题的技术栈
        
        Args:
            questions: 问题文本列表
            return_scores: 是否返回分数，默认为False
        
        Returns:
            预测结果列表，每个元素对应一个问题的预测结果
        """
        results = []
        
        for question in questions:
            if return_scores:
                techs, scores = self.predict(question, return_scores=True)
                results.append((techs, scores))
            else:
                techs = self.predict(question)
                results.append(techs)
        
        return results

# 测试代码
if __name__ == "__main__":
    # 测试预测器
    predictor = TechStackPredictor()
    
    # 准备测试问题
    test_questions = [
        "请解释Java中的多态原理和使用场景。",
        "MySQL索引的底层实现是什么？",
        "Redis缓存穿透、缓存击穿和缓存雪崩的区别是什么？",
        "如何设计一个高并发的系统？",
        "Spring框架的IoC和AOP是什么？"
    ]
    
    # 对每个测试问题进行预测并打印结果
    for question in test_questions:
        techs, scores = predictor.predict(question, return_scores=True)
        print(f"问题: {question}")
        print(f"预测技术栈: {techs}")